from langchain.text_splitter import RecursiveCharacterTextSplitter
import pandas as pd
import numpy as np
from numpy.linalg import norm
from langchain.embeddings import HuggingFaceEmbeddings



with open('./docs/西游记.txt', encoding='utf-8') as f:
    article_text = f.read()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=100,
    chunk_overlap=20,
    length_function=len,
)
texts = text_splitter.create_documents([article_text])
print(texts[1])

# 分割后的文本块还是document对象,需要提取文本字符串组成str列表
text_chunks = []
for text in texts:
    text_chunks.append(text.page_content)

# 创建一个dataframe，存储文本块等信息，便于展示
df = pd.DataFrame({'text_chunks': text_chunks})
print(df.head())

emb_model = HuggingFaceEmbeddings(model_name='./Models/m3e-base')
embeddings = emb_model.embed_documents(text_chunks)

# 将嵌入向量存储到dataframe中
df['m3e_embedding'] = embeddings
print(df.head())

# 首先需要生成用户问题的嵌入向量
users_question = "唐三藏的徒弟有哪几个？"
question_embedding = emb_model.embed_query(users_question)

# 存储问题与各个文本块的相似度
cos_sim = []
for index, row in df.iterrows():
    A = row.m3e_embedding
    B = question_embedding
    # 计算余弦相似度
    cosine = np.dot(A, B) / (norm(A) * norm(B))
    cos_sim.append(cosine)

df["cos_sim"] = cos_sim
# 按照相似度降序排列
df = df.sort_values(by=["cos_sim"], ascending=False)

print(df.head())